%%******************************************************************************************
%% Replication code: solve models
%% "Technological Rivalry and Optimal Dynamic Policy in an Open Economy"
%% by Bai, Jin, and Lu, May 2025
%%******************************************************************************************

clear all;
clc;
close all;
warning off;

irun =6;
% irun = 1, sigmaH=sigmaF=1, Ramsey
% irun = 2, sigmaH=sigmaF=2.5, Private
% irun = 3, sigmaH=sigmaF=2.5, Ramsey
% irun = 4, sigmaH=sigmaF=0, Private
% irun = 5, sigmaH=sigmaF=0, Ramsey
% irun = 6, sigma = 2.5, Markov case

if (irun==1)
    sigmaH = 1.0000001;  % risk aversion of Home
    sigmaF = 1.0000001;  % risk aversion of Foreign
elseif (irun==2 | irun==3)
    sigmaH =2.5;  % risk aversion of Home
    sigmaF = 2.5;  % risk aversion of Foreign
elseif (irun==4 | irun==5)
    sigmaH = 0;  % risk aversion of Home
    sigmaF = 0;  % risk aversion of Foreign
elseif (irun==6)
    sigmaH = 2.5;% Markov
    sigmaF = 2.5;
end


T1_frac = 0.5; % T1_frac*log(T1),starting point for T1 during the transition, fraction of Tss1
T2_frac = 1;

use_shock = 1; % 1 use alpha shock, 2 use dshock
n      =  2; % number of countries
ns     =  1; % numer of sectors
beta   = 1/(1+0.1);     % discount factor, 0.06
theta  = 2; % trade elasticity

gL      = 0.02;
delta1  = 1-1/(1+gL);
dbar   = 1.1;
rhoa   = 0.95;   % not used for transition path
std_ea = 0.001;  % not used for transition path

rhod   = 0.95;
std_ed = 0.1;

L =[1;1];
betas = ones(ns,1)/ns;

rho   = beta*(1-delta1);
rbest = delta1/(theta*(1-beta*(1-delta1))+delta1);

ncountries=n;
nsectors = ns;

alpha = [1;0.9];

if (irun<=5)

    spara(1) = ncountries;
    spara(2) = nsectors;
    spara(3) = beta;
    spara(4) = theta;
    spara(5) = sigmaH;
    spara(6) = gL;
    spara(7) = dbar;
    spara(8) = delta1;
    spara(9) = L(1);
    spara(10) = L(2);
    spara(11) = alpha(1);
    spara(12) = alpha(2);
    spara(13) = sigmaF;

    save basepara_all spara

    % run steady state
    main_notForDynare_solveSS;

    % parameters
    mpara(1) = ncountries;
    mpara(2) = nsectors;
    mpara(3) = beta;
    mpara(4) = theta;
    mpara(5) = sigmaH;
    mpara(6) = gL;
    mpara(7) = dbar;
    mpara(8) = delta1;
    mpara(9) = L(1);
    mpara(10) = L(2);
    mpara(11) = alpha(1);
    mpara(12) = alpha(2);
    mpara(13) = rhoa;
    mpara(14) = std_ea;
    mpara(15) = extax_ss;
    mpara(16) = lrss;
    mpara(17) = log(T1_frac*Tss(1));
    mpara(18) = log(T2_frac*Tss(2));

    mpara(19) = xss(1);
    mpara(20) = xss(2);
    mpara(21) = Pss(2);
    mpara(22) = wss(1);
    mpara(23) = wss(2);
    mpara(24) = rss;
    mpara(25) = pi_ss(1,2);
    mpara(26) = pi_ss(2,1);
    mpara(27) = riskfree_ss;
    mpara(28) = rmul_ss;
    mpara(29) = sigmaF;

    save setpara_ramsey mpara


    if (mod(irun,2)==0)
        % Private
        dynare private_1s2c_perfect2_varyTaux noclearall;
        cfname = ['private_sH',num2str(floor(sigmaH*100)/100),'_sF',num2str(floor(sigmaF*100)/100),'.mat'];
        copyfile('./private_1s2c_perfect2_varyTaux/Output/private_1s2c_perfect2_varyTaux_results.mat', cfname)
    else
        % Ramsey
        dynare Ramsey_1s2c_perfect2 noclearall;
        cfname = ['Ramsey_sH',num2str(floor(sigmaH*100)/100),'_sF',num2str(floor(sigmaF*100)/100),'.mat'];
        copyfile('./Ramsey_1s2c_perfect2/Output/Ramsey_1s2c_perfect2_results.mat', cfname)
    end
else
    % Markov
    gcoef0 = zeros(n-1,ns,n,ns);

    mpara(1) = n;
    mpara(2) = ns;
    mpara(3) = beta;
    mpara(4) = theta;
    mpara(5) = sigmaH;
    mpara(6) = gL;
    mpara(7) = dbar;
    mpara(8) = delta1;
    mpara(9) = rho;

    inn = 9;
    for i=1:n
        inn = inn+1;
        mpara(inn) = L(i);
    end

    for i=1:n
        for j=1:ns
            inn = inn+1;
            mpara(inn) = alpha(i,j);
        end
    end

    for j=1:ns
        inn = inn+1;
        mpara(inn) = betas(j);
    end

    inn = inn+1;
    mpara(inn) = rhoa;

    inn = inn+1;
    mpara(inn) = std_ea;

    inn = inn+1;
    mpara(inn) = rhod;

    inn = inn+1;
    mpara(inn) = std_ed;

    inn_gcoef = inn;
    for i=1:n-1
        for j=1:ns
            for ii=1:n
                for jj=1:ns
                    inn = inn+1;
                    mpara(inn) = gcoef0(i,j,ii,jj);
                end
            end
        end
    end

    mpara_stoch = mpara;
    save setpara_stoch mpara
    %% run steady state
    main_notForDynare_solveSS;

    dynare markov2 noclearall;

    %% get coefficient for dG/dT
    % simulation

    shock_matrix = randn(10000,M_.exo_nbr);
    if (use_shock==1)
        shock_matrix(:,2:M_.exo_nbr) = 0;
    else
        shock_matrix(:,1:M_.exo_nbr-1) = 0;
    end


    distg = 2;
    iterg = 0;
    maxiter = 100;
    tgcoef0 = gcoef0;
    ifrac=0.2;


    while (distg>1e-6 & iterg<maxiter)

        %dGT
        gcoef0 = tgcoef0;

        for icp=2:ncountries
            for isp=1:nsectors
                for ic=1:ncountries
                    for is=1:nsectors
                        vv=['gcoef0_',num2str(icp),num2str(isp),num2str(ic),num2str(is)];

                        set_param_value(vv,	gcoef0(icp-1,isp,ic,is));
                    end
                end
            end
        end

        save sgcoef0 gcoef0

        inn = inn_gcoef;
        for i=1:n-1
            for j=1:ns
                for ii=1:n
                    for jj=1:ns
                        inn = inn+1;
                        mpara(inn) = gcoef0(i,j,ii,jj);
                    end
                end
            end
        end

        save setpara_stoch mpara
        %
        dynare markov2 noclearall;

        %% regression
        % find GT
        vname = 'GT';
        for ic=1:ncountries
            for is=1:nsectors
                vv = [vname,num2str(ic),num2str(is)];
                mm=strcmpi(vv,M_.endo_names);
                ttt=find(mm==1);
                if (isempty(ttt)==0)
                    iGT(ic,is) = ttt; % declaration order
                else
                    iGT(ic,is)=0;
                    fprintf(1,'cannot find variable %s\n',vv);
                end
            end
        end

        vname = 'T';
        for ic=1:ncountries
            for is=1:nsectors
                vv = [vname,num2str(ic),num2str(is)];
                mm=strcmpi(vv,M_.endo_names);
                ttt=find(mm==1);
                if (isempty(ttt)==0)
                    iT(ic,is) = ttt; % declaration order
                else
                    iT(ic,is)=0;
                    fprintf(1,'cannot find variable %s\n',vv);
                end
            end
        end

        if (use_shock==1)
            vname = 'sa11';
        else
            vname = 'dshock';
        end
        vv = [vname];
        mm=strcmpi(vv,M_.endo_names);
        ttt=find(mm==1);
        if (isempty(ttt)==0)
            isa11 = ttt; % declaration order
        else
            isa11=0;
            fprintf(1,'cannot find variable %s\n',vv);
        end


        k2 = oo_.dr.state_var; % index of steady state variables

        tgcoef0 = zeros(n-1,ns,n,ns);
        ttt=find(k2==iT(1,1));
        tgcoef0(1,1,1,1) = oo_.dr.ghx(oo_.dr.inv_order_var(iGT(2,1)),ttt);%*exp(oo_.steady_state(iT(1,1)));
        ttt=find(k2==iT(2,1));
        tgcoef0(1,1,2,1) = oo_.dr.ghx(oo_.dr.inv_order_var(iGT(2,1)),ttt);%*exp(oo_.steady_state(iT(2,1)));


        % get distance
        tmp1 = (gcoef0-tgcoef0).^2;
        tmp2 = gcoef0.^2;
        distg = sqrt(sum(tmp1(:)))/(1+sqrt(sum(tmp2(:))));
        iterg = iterg+1;


        iterg
        distg
        disp('----------')
        sdistg(iterg) = distg;

    end

    save sgcoef0 gcoef0

    inn = inn_gcoef;
    for i=1:n-1
        for j=1:ns
            for ii=1:n
                for jj=1:ns
                    inn = inn+1;
                    mpara(inn) = gcoef0(i,j,ii,jj);
                end
            end
        end
    end

    mpara_stoch = mpara;
    save setpara_stoch mpara

    dynare markov2 noclearall;

    %% Markov perfect foresight

    inn = inn+1;
    mpara(inn) = T1_frac*log(Tss(1));

    inn = inn+1;
    mpara(inn) = log(Tss(2));

    mpara_perfect = mpara;

    save setpara_perfect mpara

    dynare markov2_perfect noclearall;


end


